# set up environment ----
library(data.table)
library(ggplot2)
library(Hmisc)
library(gridExtra)
source("R/00-functions.R")
load("data/analysis_data.RData")
theme_set(
  theme_bw() +
    theme(
      legend.position = "bottom",
      legend.title = element_blank(),
      plot.title = element_text(hjust = .5, size = 14),
      strip.background = element_rect(fill = NA, color = NA)))

# equivalence testing ----
X <- copy(respondent_data)
equivalence_test <- run_equiv(
  X = X[, .(
    `Age 18-24` = as.numeric(age_category == "18-24"),
    `Age 25-34` = as.numeric(age_category == "25-34"),
    `Age 35-44` = as.numeric(age_category == "35-44"),
    `Age 45-64` = as.numeric(age_category == "45-64"),
    `Age Over 65` = as.numeric(age_category == "Over 65"),
    `Did not Vote/Supp. Other Cand. in 2016` = as.numeric(support_2016 == "other/no one"),
    `Supported Clinton in 2016` = as.numeric(support_2016 == "Clinton"),
    `Supported Sanders in 2016` = as.numeric(support_2016 == "Sanders"),
    `Supported Trump in 2016` = as.numeric(support_2016 == "Trump"),
    `HS or Less` = as.numeric(education == "HS or Less"),
    `Some College` = as.numeric(education == "Some College"),
    `College Graduate` = as.numeric(education == "College Grad."),
    `Black` = as.numeric(race_ethnicity == "Black"),
    `Latino` = as.numeric(race_ethnicity == "Latino"),
    `Other/Mixed Race/Ethnicity` = as.numeric(race_ethnicity == "Other/Mixed"),
    `White` = as.numeric(race_ethnicity == "White"),
    `Female` = female,
    `Campaign Knowledge` = campaign_knowledge,
    `Follows Politics` = follows_politics,
    `Strong Democrat` = strong_democrat,
    `Trump Least Preferred` = trump_least_preferred,
    `Intend to Vote` = vote_2020,
    `Registered Voter` = registered_voter,
    Ideology = ideology,
    `Political Interest` = political_interest,
    `# Failed Attention Checks` = n_attention_fails,
    `Need for Cognition` = need_for_cognition)],
  Tr = X[, electability_last],
  w = rep(1, nrow(X)),
  epsilon.method = "strict",
  fdr_correct = TRUE)

cairo_pdf("figures/balance-full-dataset.pdf", height = 6.5, width = 12)
generate_plot(equivalence_test, display.names = equivalence_test$display.names,
  panel.widths = c(3, 1, 5, 1, 1),
  var.rounding = 2, pval.rounding = 2,
  title_text = "Equivalence Tests of Covariate Balance, Main Dataset (n = 1211)")
dev.off()

X <- copy(respondent_data[trump_least_preferred == 1 & support_2016 != "Trump" & n_attention_fails <= 1])
equivalence_test <- run_equiv(
  X = X[, .(
    `Age 18-24` = as.numeric(age_category == "18-24"),
    `Age 25-34` = as.numeric(age_category == "25-34"),
    `Age 35-44` = as.numeric(age_category == "35-44"),
    `Age 45-64` = as.numeric(age_category == "45-64"),
    `Age Over 65` = as.numeric(age_category == "Over 65"),
    `Did not Vote/Supp. Other Cand. in 2016` = as.numeric(support_2016 == "other/no one"),
    `Supported Clinton in 2016` = as.numeric(support_2016 == "Clinton"),
    `Supported Sanders in 2016` = as.numeric(support_2016 == "Sanders"),
    # `Supported Trump in 2016` = as.numeric(support_2016 == "Trump"),
    `HS or Less` = as.numeric(education == "HS or Less"),
    `Some College` = as.numeric(education == "Some College"),
    `College Graduate` = as.numeric(education == "College Grad."),
    `Black` = as.numeric(race_ethnicity == "Black"),
    `Latino` = as.numeric(race_ethnicity == "Latino"),
    `Other/Mixed Race/Ethnicity` = as.numeric(race_ethnicity == "Other/Mixed"),
    `White` = as.numeric(race_ethnicity == "White"),
    `Female` = female,
    `Campaign Knowledge` = campaign_knowledge,
    `Follows Politics` = follows_politics,
    `Strong Democrat` = strong_democrat,
    # `Trump Least Preferred` = trump_least_preferred,
    `Intend to Vote` = vote_2020,
    `Registered Voter` = registered_voter,
    Ideology = ideology,
    `Political Interest` = political_interest,
    `# Failed Attention Checks` = n_attention_fails,
    `Need for Cognition` = need_for_cognition)],
  Tr = X[, electability_last],
  w = rep(1, nrow(X)),
  epsilon.method = "strict",
  fdr_correct = TRUE)


cairo_pdf("figures/balance-main-dataset.pdf", height = 6.5, width = 12)
generate_plot(equivalence_test, display.names = equivalence_test$display.names,
  panel.widths = c(3, 1, 5, 1, 1),
  var.rounding = 2, pval.rounding = 2,
  title_text = "Equivalence Tests of Covariate Balance, Main Dataset (n = 833)")
dev.off()
